package org.anno.chunkdb.backend.query;

import com.google.common.base.Preconditions;

import it.unimi.dsi.fastutil.objects.Object2IntMap;
import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap;

import org.anno.chunkdb.Chunkdb;
import org.anno.chunkdb.Metadata;
import org.anno.chunkdb.backend.Chunk;
import org.anno.chunkdb.backend.FlyRow;
import org.anno.chunkdb.backend.Table;

import java.util.List;
import java.util.function.UnaryOperator;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * A plan for how to produce results for a query. The basic plan is {@link #fullScan}.
 */
public interface QueryPlan<R> {
  Result<R> execute();

  default QueryPlan<R> order(int column, boolean asc) {return null;}
  default QueryPlan<R> filter(int column, Chunkdb.Filter filter) {return null;}

  static QueryPlan<FlyRow> fullScan(Chunk chunk) {
     return new QueryPlan<FlyRow>() {

       public QueryPlan<FlyRow> order(int column, boolean asc) {
         return new IndexScan(chunk, column, 0, chunk.getSize(), asc);
       }

       public QueryPlan<FlyRow> filter(int column, Chunkdb.Filter filter) {
         return new IndexScan(chunk, column, 0, chunk.getSize(), true).filter(column, filter);
       }

       public Result<FlyRow> execute() {
         return new ChunkResult(chunk) {
           public Stream<FlyRow> stream() {
             return chunk.getAll();
           }
         };
       }
     };
  }

  static QueryPlan<FlyRow> from(Metadata.Table table, Chunk chunk, Chunkdb.Select query) {
    Object2IntMap<String> dict = new Object2IntOpenHashMap<>();
    for (int i = 0; i < table.getColumnCount(); i++) {
      Metadata.Column column = table.getColumn(i);
      dict.put("col" + i, i);
      dict.put(column.getName(), i);
    }

    QueryPlan<FlyRow> plan = fullScan(chunk);
    if (query.hasOrderBy()) {
      String orderBy = query.getOrderBy();
      Preconditions.checkArgument(dict.containsKey(orderBy));
      int column = dict.getInt(orderBy);
      QueryPlan<FlyRow> order = plan.order(column, query.getAscending());
      plan = order != null ? order
          : new RefinedPlan<>(plan, s -> s.sorted(chunk.getColumns().get(column).getComparator()));
    }
    for (Chunkdb.Filter filter : query.getWhereList()) {
      int c;
      if (filter.hasColumnId()) {
        c = filter.getColumnId();
      } else {
        Preconditions.checkArgument(dict.containsKey(filter.getColumn()));
        c = dict.get(filter.getColumn());
      }
      QueryPlan<FlyRow> filterPlan = plan.filter(c, filter);
      plan = filterPlan != null ? filterPlan :
          new RefinedPlan<>(plan, s -> s.filter(chunk.getColumns().get(c).getPredicate(filter)));
    }
    if (query.hasOffset()) {
      plan = new RefinedPlan<>(plan, s -> s.skip(query.getOffset()));
    }
    if (query.hasLimit()) {
      plan = new RefinedPlan<>(plan, s -> s.limit(query.getLimit()));
    }
    return plan;
  }


  class RefinedPlan<R> implements QueryPlan<R> {
    private final QueryPlan<R> delegate;
    private final UnaryOperator<Stream<R>> operator;

    protected RefinedPlan(QueryPlan<R> delegate, UnaryOperator<Stream<R>> operator) {
      this.delegate = delegate;
      this.operator = operator;
    }

    public Result<R> execute() {
      final Result<R> result = delegate.execute();
      return new Result<R>() {
        public Stream<R> stream() {
          return operator.apply(result.stream());
        }
        public List<ResultColumn<R>> columns() {
          return result.columns();
        }
      };
    }

    public QueryPlan<R> order(int column, boolean asc) {
      return delegate.order(column, asc);
    }

    public QueryPlan<R> filter(int column, Chunkdb.Filter filter) {
      return delegate.filter(column, filter);
    }
  }

  class IndexScan implements QueryPlan<FlyRow> {
    private final Chunk chunk;
    private final int column;
    private final int start;
    private final int end;
    private final boolean asc;

    private IndexScan(Chunk chunk, int column, int start, int end, boolean asc) {
      this.chunk = chunk;
      this.column = column;
      this.start = start;
      this.end = end;
      this.asc = asc;
    }

    public Result<FlyRow> execute() {
      return new ChunkResult(chunk) {
        public Stream<FlyRow> stream() {
          return getColumn().getRowRange(start, end, asc).mapToObj(chunk.getIntToRowFunction());
        }
      };
    }

    public QueryPlan<FlyRow> order(int column, boolean asc) {
      return column != this.column ? null
          : asc == this.asc ? this
              : new IndexScan(chunk, column, start, end, asc);
    }

    public QueryPlan<FlyRow> filter(int column, Chunkdb.Filter filter) {
      return column != this.column ? null
          : new IndexScan(chunk, column, getStart(filter), getEnd(filter), asc);
    }

    private int getStart(Chunkdb.Filter filter) {
      switch (filter.getOp()) {
        case EQ:
        case GE:
          return Math.max(start, filter.hasStringRhs()
              ? getColumn().getGERank(filter.getStringRhsBytes().toByteArray())
              : getColumn().getGERank(getRhsAsDouble(filter)));
        case GT:
          return Math.max(start, filter.hasStringRhs()
              ? getColumn().getGTRank(filter.getStringRhsBytes().toByteArray())
              : getColumn().getGTRank(getRhsAsDouble(filter)));
      }
      return start;
    }

    private Table.Column<FlyRow> getColumn() {
      return chunk.getColumns().get(column);
    }

    private int getEnd(Chunkdb.Filter filter) {
      switch (filter.getOp()) {
        case LT:
          return Math.min(end, filter.hasStringRhs()
              ? getColumn().getLTRank(filter.getStringRhsBytes().toByteArray())
              : getColumn().getLTRank(getRhsAsDouble(filter)));
        case LE:
        case EQ:
          return Math.min(end, filter.hasStringRhs()
              ? getColumn().getLERank(filter.getStringRhsBytes().toByteArray())
              : getColumn().getLERank(getRhsAsDouble(filter)));

      }
      return end;
    }

    private double getRhsAsDouble(Chunkdb.Filter filter) {
      return filter.hasDoubleRhs() ? filter.getDoubleRhs() : (double) filter.getLongRhs();
    }
  }

  abstract class ChunkResult implements Result<FlyRow> {

    protected final Chunk chunk;

    public ChunkResult(Chunk chunk) {
      this.chunk = chunk;
    }

    @SuppressWarnings("unchecked")
    public List<ResultColumn<FlyRow>> columns() {
      return chunk.getColumns().stream()
          .map(c -> (ResultColumn<FlyRow>) c)
          .collect(Collectors.toList());
    }
  }
}
